import io
import os
import tempfile
import uuid
from datetime import datetime

import cv2
import numpy as np
from PIL import Image
from flask import Blueprint, request, send_file, Response
from ultralytics import YOLO
from werkzeug.utils import secure_filename

# 创建一个蓝图对象
bp = Blueprint('bp', __name__)

# 加载训练后的模型
model = YOLO('weight/best.pt')
# 预热一下，避免第一次请求缓慢超时
model.predict()


# 计算图片4个顶点坐标
def calculate_rectangle_vertices(x_center, y_center, width, height):
    left_top = (x_center - width / 2, y_center - height / 2)
    right_top = (x_center + width / 2, y_center - height / 2)
    right_bottom = (x_center + width / 2, y_center + height / 2)
    left_bottom = (x_center - width / 2, y_center + height / 2)

    # 返回包含四个顶点坐标的列表（每个坐标是元组）
    return [left_top, right_top, right_bottom, left_bottom]


# 裁切图片
class Crop:
    def handler(self, files: request, conf: float):
        if 'file' not in files:
            raise Exception("请上传文件")

        file = files['file']
        # 如果用户没有选择文件，浏览器也会提交一个没有文件名的空部分
        if file is None or file.filename == '':
            raise Exception("请上传文件")

        image = np.frombuffer(file.read(), np.uint8)
        image = cv2.imdecode(image, cv2.IMREAD_COLOR)
        # 推理
        results = model.predict(image, conf=conf)
        # 对结果进行裁切
        img_byte_arr = self.__handler2(results)

        return img_byte_arr

    # 从结果中获取矩形顶点坐标，并提取图片，返回图片路径
    def __handler2(self, results):
        for r in results:  # 对结果进行处理
            if (r.boxes.conf is None) or (len(r.boxes.conf) == 0):
                raise Exception("没有检测到物体")

            for detection in r.boxes.xywh:  # detection获取边界信息并保留在列表中
                # 提取中心点 x 坐标、中心点 y 坐标、宽度和高度
                x = detection[0].item()  # x坐标为列表中的第一个
                y = detection[1].item()  # y坐标为列表中的第二个
                w = detection[2].item()  # 矩形框宽度为列表中的第三个
                h = detection[3].item()  # 矩形框高度为列表中的第四个

                # 获取4个坐标点
                vertices = calculate_rectangle_vertices(x, y, w, h)

                # 裁切出图片
                x1, y1 = vertices[0]  # 左上角坐标
                x2, y2 = vertices[2]  # 右下角坐标
                img_cv2 = cv2.cvtColor(r.orig_img, cv2.COLOR_BGR2RGB)
                img_pil = Image.fromarray(img_cv2)
                box = (x1, y1, x2, y2)  # 裁剪大小
                new_img = img_pil.crop(box)

                # 返回字节流
                img_byte_arr = io.BytesIO()
                new_img.save(img_byte_arr, format='JPEG')
                img_byte_arr = img_byte_arr.getvalue()

                return img_byte_arr


# 接受文件上传，调用模型处理
@bp.route("/uploadFile", methods=["POST"])
def upload_file():
    crop = Crop()
    # 返回字节流
    byte_arr = crop.handler(request.files, 0.9)

    response = Response(byte_arr, mimetype='image/jpeg')

    # 设置HTTP响应头，以便浏览器正确处理
    response.headers['Content-Disposition'] = 'inline; filename="image.jpg"'

    return response
